#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/3/8 20:00
# @Author  : jeff
# @FileName: user_model
import sqlalchemy
from sqlalchemy.orm import declarative_base, relationship
from sqlalchemy.sql import expression

from src.model.base_model import BaseModel, db

'''
  `id` bigint unsigned NOT NULL AUTO_INCREMENT,
  `uuid` varchar(37) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci DEFAULT NULL COMMENT '子任务uuid',
  `work_uuid` varchar(37) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci NOT NULL COMMENT '总任务uuid',
  `params` json DEFAULT NULL COMMENT '任务参数',
  `status` varchar(32) CHARACTER SET utf8mb4 COLLATE utf8mb4_general_ci NOT NULL DEFAULT 'pending' COMMENT '任务状态',
  `retry` tinyint unsigned NOT NULL DEFAULT '0' COMMENT '重试次数',
  `priority` tinyint unsigned NOT NULL DEFAULT '0' COMMENT '优先级',
  `execute_time` json DEFAULT NULL COMMENT '执行时间',
  `progress` tinyint unsigned NOT NULL DEFAULT '0' COMMENT '进度',
  `create_time` datetime(3) NOT NULL,
  `update_time` datetime(3) NOT NULL,
'''


class ScanTask(BaseModel):
    __tablename__ = 'port_scan_task'

    id = db.Column(db.Integer, primary_key=True, autoincrement=True)
    uuid = db.Column(db.String(37), nullable=False, unique=True)
    work_uuid = db.Column(db.String(37), nullable=False)
    params = db.Column(db.JSON(), nullable=False, default=lambda: {})
    status = db.Column(db.String(32))
    retry = db.Column(db.Integer)
    priority = db.Column(db.Integer)
    execute_time = db.Column(sqlalchemy.DateTime, server_default=sqlalchemy.func.current_timestamp())
    progress = db.Column(db.Integer)
    create_time = db.Column(sqlalchemy.DateTime, server_default=sqlalchemy.func.current_timestamp())
    update_time = db.Column(sqlalchemy.DateTime, server_default=sqlalchemy.func.current_timestamp(),
                            onupdate=sqlalchemy.func.current_timestamp())

    @classmethod
    def get_pending_task(cls, status, uuid):
        try:
            db.session.expire_all()  # 让 sessio
            pending_task = cls.query.with_for_update().filter(ScanTask.status == status,
                                                             ScanTask.uuid == uuid).order_by(
                ScanTask.priority).first()
            return pending_task
        except Exception:
            raise

    @classmethod
    def get_status_task(cls, status):
        try:
            db.session.expire_all()  # 让 sessio
            pending_task = cls.query.with_for_update().filter(ScanTask.status == status).order_by(
                ScanTask.priority).first()
            return pending_task
        except Exception:
            raise

    @classmethod
    def query_by_uuid(cls,uuid):
        try:
            db.session.expire_all()  # 让 session 失效，下一次查询会从数据库获取最新数据
            pending_task = cls.query.with_for_update().filter(ScanTask.uuid == uuid).order_by(
                ScanTask.priority).first()
            return pending_task
        except Exception:
            raise

    @classmethod
    def get_work_uuid_task(cls, work_uuid):
        try:
            db.session.expire_all()  # 让 sessio
            pending_task = cls.query.with_for_update().filter(ScanTask.work_uuid == work_uuid)
            return pending_task
        except Exception:
            raise

    @classmethod
    def get_work_uuid_task_status(cls, work_uuid, status):
        try:
            db.session.expire_all()  # 让 sessio
            pending_task = cls.query.with_for_update().filter(ScanTask.work_uuid == work_uuid,
                                                             ScanTask.status == status).order_by(
                ScanTask.priority).first()
            return pending_task
        except Exception:
            raise
